{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook written by [Zhedong Zheng](https://github.com/zhedongzheng)\n",
    "\n",
    "![title](img/seq2seq.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "PARAMS = {\n",
    "    'embed_dims': 15,\n",
    "    'rnn_size': 50,\n",
    "    'num_layers': 1,\n",
    "    'beam_width': 5,\n",
    "    'clip_norm': 5.0,\n",
    "    'batch_size': 128,\n",
    "    'n_epochs': 60,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_data(path):\n",
    "    with open(path, 'r', encoding='utf-8') as f:\n",
    "        return f.read()\n",
    "\n",
    "    \n",
    "def build_map(data):\n",
    "    specials = ['<PAD>', '<GO>',  '<EOS>', '<UNK>']\n",
    "    chars = list(set([char for line in data.split('\\n') for char in line]))\n",
    "    idx2char = {idx: char for idx, char in enumerate(specials + chars)}\n",
    "    char2idx = {char: idx for idx, char in idx2char.items()}\n",
    "    return idx2char, char2idx\n",
    "\n",
    "\n",
    "def preprocess_data():\n",
    "    source = read_data('../temp/letters_source.txt')\n",
    "    target = read_data('../temp/letters_target.txt')\n",
    "\n",
    "    PARAMS['src_idx2char'], PARAMS['src_char2idx'] = build_map(source)\n",
    "    PARAMS['tgt_idx2char'], PARAMS['tgt_char2idx'] = build_map(target)\n",
    "\n",
    "    src_idx = [[PARAMS['src_char2idx'].get(char, 3) for char in line] for line in source.split('\\n')]\n",
    "    tgt_idx = [[PARAMS['tgt_char2idx'].get(char, 3) for char in line]+[2] for line in target.split('\\n')]\n",
    "\n",
    "    return src_idx, tgt_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_sent_batch(sent_batch):\n",
    "    max_sent_len = max([len(sent) for sent in sent_batch])\n",
    "    padded_seqs = [(sent + [0]*(max_sent_len - len(sent))) for sent in sent_batch]\n",
    "    return padded_seqs\n",
    "\n",
    "\n",
    "def next_train_batch(src_idx, tgt_idx):\n",
    "    for i in range(0, len(src_idx), PARAMS['batch_size']):\n",
    "        padded_src = pad_sent_batch(src_idx[i: i+PARAMS['batch_size']])\n",
    "        padded_tgt = pad_sent_batch(tgt_idx[i: i+PARAMS['batch_size']])\n",
    "        yield padded_src, padded_tgt\n",
    "\n",
    "        \n",
    "def train_input_fn(src_idx, tgt_idx):\n",
    "    dataset = tf.data.Dataset.from_generator(\n",
    "        lambda: next_train_batch(src_idx, tgt_idx),\n",
    "        (tf.int32, tf.int32),\n",
    "        (tf.TensorShape([None, None]), tf.TensorShape([None, None])))\n",
    "    dataset = dataset.repeat(PARAMS['n_epochs'])\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    return iterator.get_next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clip_grads(loss):\n",
    "    variables = tf.trainable_variables()\n",
    "    grads = tf.gradients(loss, variables)\n",
    "    clipped_grads, _ = tf.clip_by_global_norm(grads, PARAMS['clip_norm'])\n",
    "    return zip(clipped_grads, variables)\n",
    "\n",
    "\n",
    "def rnn_cell():\n",
    "    def cell_fn():\n",
    "        cell = tf.nn.rnn_cell.GRUCell(PARAMS['rnn_size'],\n",
    "                                      kernel_initializer=tf.orthogonal_initializer())\n",
    "        return cell\n",
    "    return tf.nn.rnn_cell.MultiRNNCell([cell_fn() for _ in range(PARAMS['num_layers'])])\n",
    "\n",
    "\n",
    "def dec_cell(enc_out, enc_seq_len):\n",
    "    attention = tf.contrib.seq2seq.BahdanauAttention(\n",
    "        num_units = PARAMS['rnn_size'],\n",
    "        memory = enc_out,\n",
    "        memory_sequence_length = enc_seq_len)\n",
    "    \n",
    "    return tf.contrib.seq2seq.AttentionWrapper(\n",
    "        cell = rnn_cell(),\n",
    "        attention_mechanism = attention,\n",
    "        attention_layer_size = PARAMS['rnn_size'])\n",
    "\n",
    "\n",
    "def dec_input(labels):\n",
    "    x = tf.fill([tf.shape(labels)[0], 1], PARAMS['tgt_char2idx']['<GO>'])\n",
    "    x = tf.to_int32(x)\n",
    "    return tf.concat([x, labels[:, :-1]], 1)\n",
    "\n",
    "\n",
    "def forward(inputs, labels, mode):\n",
    "    is_training = (mode == tf.estimator.ModeKeys.TRAIN)\n",
    "    enc_seq_len = tf.count_nonzero(inputs, 1, dtype=tf.int32)\n",
    "    batch_sz = tf.shape(inputs)[0]\n",
    "    \n",
    "    with tf.variable_scope('Encoder'):\n",
    "        embedding = tf.get_variable('lookup_table',\n",
    "                                    [len(PARAMS['src_char2idx']), PARAMS['embed_dims']])\n",
    "        x = tf.nn.embedding_lookup(embedding, inputs)\n",
    "        enc_out, enc_state = tf.nn.dynamic_rnn(rnn_cell(), x, enc_seq_len, dtype=tf.float32)\n",
    "        \n",
    "    with tf.variable_scope('Decoder'):\n",
    "        output_proj = tf.layers.Dense(len(PARAMS['tgt_char2idx']))\n",
    "        \n",
    "        if is_training:\n",
    "            cell = dec_cell(enc_out, enc_seq_len)\n",
    "            dec_seq_len = tf.count_nonzero(labels, 1, dtype=tf.int32)\n",
    "            \n",
    "            init_state = cell.zero_state(batch_sz, tf.float32).clone(\n",
    "                cell_state=enc_state)\n",
    "            \n",
    "            helper = tf.contrib.seq2seq.TrainingHelper(\n",
    "                inputs = tf.nn.embedding_lookup(embedding, dec_input(labels)),\n",
    "                sequence_length = dec_seq_len)\n",
    "            decoder = tf.contrib.seq2seq.BasicDecoder(\n",
    "                cell = cell,\n",
    "                helper = helper,\n",
    "                initial_state = init_state,\n",
    "                output_layer = output_proj)\n",
    "            decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "                decoder = decoder,\n",
    "                maximum_iterations = tf.reduce_max(dec_seq_len))\n",
    "            \n",
    "            return decoder_output.rnn_output\n",
    "        else:\n",
    "            enc_out_t = tf.contrib.seq2seq.tile_batch(enc_out, PARAMS['beam_width'])\n",
    "            enc_state_t = tf.contrib.seq2seq.tile_batch(enc_state, PARAMS['beam_width'])\n",
    "            enc_seq_len_t = tf.contrib.seq2seq.tile_batch(enc_seq_len, PARAMS['beam_width'])\n",
    "            \n",
    "            cell = dec_cell(enc_out_t, enc_seq_len_t)\n",
    "            \n",
    "            init_state = cell.zero_state(batch_sz*PARAMS['beam_width'], tf.float32).clone(\n",
    "                cell_state=enc_state_t)\n",
    "            \n",
    "            decoder = tf.contrib.seq2seq.BeamSearchDecoder(\n",
    "                cell = cell,\n",
    "                embedding = embedding,\n",
    "                start_tokens = tf.tile(tf.constant([PARAMS['tgt_char2idx']['<GO>']], tf.int32),\n",
    "                                       [batch_sz]),\n",
    "                end_token = PARAMS['tgt_char2idx']['<EOS>'],\n",
    "                initial_state = init_state,\n",
    "                beam_width = PARAMS['beam_width'],\n",
    "                output_layer = output_proj)\n",
    "            decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "                decoder = decoder)\n",
    "            \n",
    "            return decoder_output.predicted_ids[:, :, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_fn(features, labels, mode):\n",
    "    logits_or_ids = forward(features, labels, mode)\n",
    "    \n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        return tf.estimator.EstimatorSpec(mode, predictions=logits_or_ids)\n",
    "        \n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        loss_op = tf.contrib.seq2seq.sequence_loss(logits = logits_or_ids,\n",
    "                                                   targets = labels,\n",
    "                                                   weights = tf.to_float(tf.sign(labels)))\n",
    "        train_op = tf.train.AdamOptimizer().apply_gradients(\n",
    "            clip_grads(loss_op),\n",
    "            global_step = tf.train.get_global_step())\n",
    "        \n",
    "        return tf.estimator.EstimatorSpec(\n",
    "            mode=mode, loss=loss_op, train_op=train_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using default config.\n",
      "WARNING:tensorflow:Using temporary folder as model directory: /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpz4kwntd1\n",
      "INFO:tensorflow:Using config: {'_model_dir': '/var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpz4kwntd1', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': None, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x1141e1cf8>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Create CheckpointSaverHook.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Saving checkpoints for 1 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpz4kwntd1/model.ckpt.\n",
      "INFO:tensorflow:loss = 3.4004636, step = 1\n",
      "INFO:tensorflow:global_step/sec: 29.5445\n",
      "INFO:tensorflow:loss = 2.007127, step = 101 (3.386 sec)\n",
      "INFO:tensorflow:global_step/sec: 40.9444\n",
      "INFO:tensorflow:loss = 1.253626, step = 201 (2.442 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.6942\n",
      "INFO:tensorflow:loss = 0.952549, step = 301 (2.399 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.5978\n",
      "INFO:tensorflow:loss = 0.6709183, step = 401 (2.404 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.4552\n",
      "INFO:tensorflow:loss = 0.54064995, step = 501 (2.412 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.6015\n",
      "INFO:tensorflow:loss = 0.3710653, step = 601 (2.404 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.5364\n",
      "INFO:tensorflow:loss = 0.21596724, step = 701 (2.408 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.8151\n",
      "INFO:tensorflow:loss = 0.16904877, step = 801 (2.391 sec)\n",
      "INFO:tensorflow:global_step/sec: 40.3285\n",
      "INFO:tensorflow:loss = 0.12318576, step = 901 (2.480 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.5257\n",
      "INFO:tensorflow:loss = 0.07829217, step = 1001 (2.408 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.6974\n",
      "INFO:tensorflow:loss = 0.054099064, step = 1101 (2.398 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.4775\n",
      "INFO:tensorflow:loss = 0.03172837, step = 1201 (2.411 sec)\n",
      "INFO:tensorflow:global_step/sec: 40.8116\n",
      "INFO:tensorflow:loss = 0.032596204, step = 1301 (2.450 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.3138\n",
      "INFO:tensorflow:loss = 0.026931563, step = 1401 (2.420 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.3064\n",
      "INFO:tensorflow:loss = 0.009059685, step = 1501 (2.421 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.0406\n",
      "INFO:tensorflow:loss = 0.015354849, step = 1601 (2.437 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.7992\n",
      "INFO:tensorflow:loss = 0.02121437, step = 1701 (2.392 sec)\n",
      "INFO:tensorflow:global_step/sec: 40.7545\n",
      "INFO:tensorflow:loss = 0.014728889, step = 1801 (2.454 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.7608\n",
      "INFO:tensorflow:loss = 0.040140267, step = 1901 (2.395 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.5718\n",
      "INFO:tensorflow:loss = 0.009654822, step = 2001 (2.405 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.7253\n",
      "INFO:tensorflow:loss = 0.008483502, step = 2101 (2.397 sec)\n",
      "INFO:tensorflow:global_step/sec: 40.8904\n",
      "INFO:tensorflow:loss = 0.005563527, step = 2201 (2.446 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.8282\n",
      "INFO:tensorflow:loss = 0.004427869, step = 2301 (2.391 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.4124\n",
      "INFO:tensorflow:loss = 0.007250657, step = 2401 (2.415 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.7822\n",
      "INFO:tensorflow:loss = 0.0044385144, step = 2501 (2.393 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.575\n",
      "INFO:tensorflow:loss = 0.0029519787, step = 2601 (2.406 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.4879\n",
      "INFO:tensorflow:loss = 0.0029069576, step = 2701 (2.410 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.4444\n",
      "INFO:tensorflow:loss = 0.0021211049, step = 2801 (2.413 sec)\n",
      "INFO:tensorflow:global_step/sec: 40.5312\n",
      "INFO:tensorflow:loss = 0.002023036, step = 2901 (2.467 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.4168\n",
      "INFO:tensorflow:loss = 0.0018967196, step = 3001 (2.414 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.5463\n",
      "INFO:tensorflow:loss = 0.0019105093, step = 3101 (2.407 sec)\n",
      "INFO:tensorflow:global_step/sec: 40.7777\n",
      "INFO:tensorflow:loss = 0.0029551133, step = 3201 (2.452 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.533\n",
      "INFO:tensorflow:loss = 0.0018048794, step = 3301 (2.408 sec)\n",
      "INFO:tensorflow:global_step/sec: 40.5773\n",
      "INFO:tensorflow:loss = 0.0016996324, step = 3401 (2.464 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.6411\n",
      "INFO:tensorflow:loss = 0.0010506083, step = 3501 (2.401 sec)\n",
      "INFO:tensorflow:global_step/sec: 39.9539\n",
      "INFO:tensorflow:loss = 0.0009840138, step = 3601 (2.503 sec)\n",
      "INFO:tensorflow:global_step/sec: 39.8003\n",
      "INFO:tensorflow:loss = 0.0009346631, step = 3701 (2.512 sec)\n",
      "INFO:tensorflow:global_step/sec: 40.0875\n",
      "INFO:tensorflow:loss = 0.0011153592, step = 3801 (2.494 sec)\n",
      "INFO:tensorflow:global_step/sec: 39.6236\n",
      "INFO:tensorflow:loss = 0.0009042719, step = 3901 (2.524 sec)\n",
      "INFO:tensorflow:global_step/sec: 39.8449\n",
      "INFO:tensorflow:loss = 0.0010135552, step = 4001 (2.510 sec)\n",
      "INFO:tensorflow:global_step/sec: 40.4625\n",
      "INFO:tensorflow:loss = 0.00070887926, step = 4101 (2.471 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.9838\n",
      "INFO:tensorflow:loss = 0.0007847747, step = 4201 (2.382 sec)\n",
      "INFO:tensorflow:global_step/sec: 41.6519\n",
      "INFO:tensorflow:loss = 0.00066128134, step = 4301 (2.401 sec)\n",
      "INFO:tensorflow:global_step/sec: 39.3455\n",
      "INFO:tensorflow:loss = 0.000559325, step = 4401 (2.542 sec)\n",
      "INFO:tensorflow:global_step/sec: 35.4831\n",
      "INFO:tensorflow:loss = 0.0005396663, step = 4501 (2.818 sec)\n",
      "INFO:tensorflow:global_step/sec: 37.8901\n",
      "INFO:tensorflow:loss = 0.00035508416, step = 4601 (2.640 sec)\n",
      "INFO:tensorflow:global_step/sec: 37.7436\n",
      "INFO:tensorflow:loss = 0.00027019295, step = 4701 (2.649 sec)\n",
      "INFO:tensorflow:Saving checkpoints for 4740 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpz4kwntd1/model.ckpt.\n",
      "INFO:tensorflow:Loss for final step: 0.00010065189.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpz4kwntd1/model.ckpt-4740\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "\n",
      "IN: apple\n",
      "OUT: a e l p p <EOS> <EOS> <EOS>\n",
      "\n",
      "IN: common\n",
      "OUT: c m m n o o <EOS> <EOS>\n",
      "\n",
      "IN: zhedong\n",
      "OUT: d e g h n o z <EOS>\n"
     ]
    }
   ],
   "source": [
    "def infe_inps(str_li):\n",
    "    max_len = max([len(s) for s in str_li])\n",
    "    xs = [[PARAMS['src_char2idx'].get(c, 3) for c in s] for s in str_li]\n",
    "    return tf.keras.preprocessing.sequence.pad_sequences(xs, max_len, padding='post')\n",
    "\n",
    "\n",
    "def demo(xs, preds):\n",
    "    for x, pred in zip(xs, preds):\n",
    "        print('\\nIN: {}'.format(x))\n",
    "        print('OUT: {}'.format(' '.join([PARAMS['tgt_idx2char'][i] for i in pred])))\n",
    "    \n",
    "\n",
    "def main():\n",
    "    src_idx, tgt_idx = preprocess_data()\n",
    "    \n",
    "    estimator = tf.estimator.Estimator(model_fn)\n",
    "    \n",
    "    estimator.train(lambda: train_input_fn(src_idx, tgt_idx))\n",
    "    \n",
    "    xs = ['apple', 'common', 'zhedong']\n",
    "    \n",
    "    preds = list(estimator.predict(tf.estimator.inputs.numpy_input_fn(\n",
    "        x = infe_inps(xs),\n",
    "        shuffle = False)))\n",
    "    \n",
    "    demo(xs, preds)\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
